import torch
import itertools
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks
from . import EnDecodes
from . import contraloss
import os


class HADISGModel(BaseModel):
    """
    hadi model which has the same content encoder and different decoder
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
        A (source domain), B (target domain).
        Generators: G_A: A -> B; G_B: B -> A.
        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
        Dropout is not used in the original CycleGAN paper.
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.add_argument('--lambda_MSE', type=float, default=10.0, help='weight')
        parser.add_argument('--lambda_GAN_Image', type=float, default=0.1, help='weight')
        parser.add_argument('--lambda_GAN_Feature', type=float, default=0.01, help='weight')
        parser.add_argument('--lambda_Contra', type=float, default=0.1, help='weight')
        parser.add_argument('--lambda_Consis', type=float, default=10.0, help='weight')
        parser.add_argument('--lambda_KL', type=float, default=0.1, help='weight')
        parser.add_argument('--feaLength', type=int, default=64, help='weight')
        parser.add_argument('--pre_path', type=str, help='pretained path')
        parser.add_argument('--init_derain', type=str, default='1,3', help='low which pred-trained model. 0 for no pre-train, 1 for loading G1 pre-trained parameters, 3 for loading G3')
        return parser

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        self.loss_names = ['MSE','Contra','GAN_Image','Consis','D_B','KL']
        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        self.visual_names = ['Os','Bs','pred_Bs','pred_Rs','Or','pred_Br','pred_Rr']
        # self.visual_names = ['Or','pred_Br','pred_Rr']

        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        if self.isTrain:
            self.model_names = ['E_B','E_Rc','E_Rs','E_Rr',
                                'G_Bs','G_Br','G_R','D_B','F']
        else:  # during test time, only load Gs
            self.model_names = ['E_B','E_Rc','E_Rs','E_Rr',
                                'G_Bs','G_Br','G_R','D_B','F']

        # define networks (both Generators and discriminators)
        # The naming is different from those used in the paper.
        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
        self.feaLength = opt.feaLength

        self.netE_B = EnDecodes.define_Encoder('encoder_content', opt.input_nc, opt.feaLength, self.gpu_ids)
        self.netE_Rc = EnDecodes.define_Encoder('encoder_content', opt.input_nc, opt.feaLength, self.gpu_ids)
        self.netE_Rs = EnDecodes.define_Encoder('encoder_content', opt.input_nc, opt.feaLength, self.gpu_ids)
        self.netE_Rr = EnDecodes.define_Encoder('encoder_content', opt.input_nc, opt.feaLength, self.gpu_ids)

        self.netG_Bs = EnDecodes.define_Encoder('decoder_single', opt.feaLength, 3, self.gpu_ids)
        self.netG_Br = EnDecodes.define_Encoder('decoder_single', opt.feaLength, 3, self.gpu_ids)
        self.netG_R = EnDecodes.define_Encoder('decoder_double', opt.feaLength, 3, self.gpu_ids)

        self.netD_B = networks.define_D(opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type,
                                        opt.init_gain, self.gpu_ids)
        self.netF = EnDecodes.define_Encoder('mlp', opt.feaLength, 64, self.gpu_ids)

        if self.isTrain:  # define discriminators
            if self.opt.init_derain != '0':
                if self.opt.pre_path == None:
                    print('FileNotFoundError: PreTrained_load_path is not found!')
                    raise FileNotFoundError
                E_B_file = 'latest_net_E_Bs.pth'
                G_B_file = 'latest_net_G_B.pth'
                state_dict_E = torch.load(os.path.join(self.opt.pre_path, E_B_file), map_location=str(self.device))
                state_dict_G = torch.load(os.path.join(self.opt.pre_path, G_B_file), map_location=str(self.device))
                model_E_B = self.netE_B.module
                model_E_B.load_state_dict(state_dict_E)
                print('initial with %s successfully!' % self.opt.pre_path)
                if '1' in self.opt.init_derain:
                    model_G_Bs = self.netG_Bs.module
                    model_G_Bs.load_state_dict(state_dict_G)
                    print('initial E_Bswith %s successfully!' % self.opt.pre_path)
                if '3' in self.opt.init_derain:
                    model_G_Br = self.netG_Br.module
                    model_G_Br.load_state_dict(state_dict_G)
                    print('initial netG3 with %s successfully!' % self.opt.pre_path)


            self.Bs_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.Br_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fBs_pool = ImagePool(opt.pool_size)
            self.fBr_pool = ImagePool(opt.pool_size)
            # define loss functions
            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
            self.criterionCycle = torch.nn.L1Loss()
            self.criterionIdt = torch.nn.L1Loss()
            self.criterion_MSE = torch.nn.MSELoss()
            self.criterion_Contra = contraloss.SupConLoss()
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            if '1' in self.opt.init_derain:
                parameters_list = [dict(params=self.netG_Bs.parameters(), lr=opt.lr / 100)]
            else:
                parameters_list = [dict(params=self.netG_Bs.parameters(), lr=opt.lr)]
            if '3' in self.opt.init_derain:
                parameters_list.append(dict(params=self.netG_Br.parameters(), lr=opt.lr / 10))
            else:
                parameters_list.append(dict(params=self.netG_Br.parameters(), lr=opt.lr))
            parameters_list.append(dict(params=self.netE_Rs.parameters(), lr=opt.lr))
            parameters_list.append(dict(params=self.netE_Rr.parameters(), lr=opt.lr))
            parameters_list.append(dict(params=self.netE_Rc.parameters(), lr=opt.lr))
            parameters_list.append(dict(params=self.netE_B.parameters(), lr=opt.lr))
            parameters_list.append(dict(params=self.netG_R.parameters(), lr=opt.lr))
            self.optimizer_G = torch.optim.Adam(parameters_list, lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_D_B = torch.optim.Adam(self.netD_B.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D_B)

    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        # AtoB = self.opt.direction == 'AtoB'
        # self.real_A = input['A' if AtoB else 'B'].to(self.device)
        # self.real_B = input['B' if AtoB else 'A'].to(self.device)
        # self.image_paths = input['A_paths' if AtoB else 'B_paths']

        self.Os = input['O_s'].to(self.device)
        self.Bs = input['B_s'].to(self.device)
        self.Or = input['O_t'].to(self.device)
        # if not self.isTrain:
        if self.opt.Bt_access:
            self.Bt = input['B_t'].to(self.device)
        # self.Bt = input['B_t'].to(self.device)
        self.image_paths = input['path']  # for test

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        # Image Space
        self.f_Bs = self.netE_B(self.Os)
        self.f_Br = self.netE_B(self.Or)
        self.pred_Bs = self.netG_Bs(self.f_Bs)
        self.pred_Br = self.netG_Br(self.f_Br)

        # Rain Space
        self.f_Rss = self.netE_Rs(self.Os)  # Rs special
        self.f_Rsc = self.netE_Rc(self.Os)  # Rs common
        self.f_Rrs = self.netE_Rr(self.Or)  # Rr special
        self.f_Rrc = self.netE_Rc(self.Or)  # Rr common
        self.pred_Rs = self.netG_R(self.f_Rsc, self.f_Rss)
        self.pred_Rr = self.netG_R(self.f_Rrc, self.f_Rrs)

        self.recon_Os = self.pred_Bs + self.pred_Rs
        self.recon_Or = self.pred_Br + self.pred_Rr

        # Reconstruction
        self.recon_f_Bs = self.netE_B(self.recon_Os)
        self.recon_f_Br = self.netE_B(self.recon_Or)


    def cal_GAN_loss_D_basic(self, netD, real, fake):
        pred_real = netD(real)
        loss_D_real = self.criterionGAN(pred_real, True)
        # Fake
        pred_fake = netD(fake.detach())
        loss_D_fake = self.criterionGAN(pred_fake, False)
        # Combined loss and calculate gradients
        loss_D = (loss_D_real + loss_D_fake) * 0.5
        return loss_D

    def backward_D_B(self):
        pred_Br = self.Br_pool.query(self.pred_Br)
        pred_Bs = self.Bs_pool.query(self.pred_Bs)
        self.loss_D_B = self.cal_GAN_loss_D_basic(self.netD_B, pred_Bs, pred_Br)
        self.loss_D_B.backward()

    def backward_G(self):
        lambda_MSE = self.opt.lambda_MSE
        lambda_GAN_Image = self.opt.lambda_GAN_Image
        lambda_Contra = self.opt.lambda_Contra
        lambda_Consis = self.opt.lambda_Consis
        lambda_KL = self.opt.lambda_KL

        # Synthetic MSE loss
        self.MSE_B = self.criterion_MSE(self.Bs, self.pred_Bs)
        self.MSE_R = self.criterion_MSE(self.Os - self.Bs, self.pred_Rs)
        self.loss_MSE = self.MSE_B + self.MSE_R

        # Image Space GAN loss
        self.loss_GAN_Image = self.criterionGAN(self.netD_B(self.pred_Br), True)
        # self.GAN_Bs = self.criterionGAN(self.netD_B(self.pred_Bs), False)
        self.loss_GAN = lambda_GAN_Image * self.loss_GAN_Image

        # Rain Space Contrastive loss
        fl_Rss = self.f_Rss.flatten(1,3) # b
        fl_Rsc = self.f_Rsc.flatten(1,3)
        fl_Rrs = self.f_Rrs.flatten(1,3)
        fl_Rrc = self.f_Rrc.flatten(1,3)
        f_Rss = self.netF(fl_Rss)
        f_Rsc = self.netF(fl_Rsc)
        f_Rrs = self.netF(fl_Rrs)
        f_Rrc = self.netF(fl_Rrc)
        labels = torch.tensor([1,2,3,2]).cuda()
        features = torch.cat([f_Rss,f_Rsc,f_Rrs,f_Rrc], dim=0).unsqueeze(1)
        features = torch.cat([features,features], dim=1)
        self.loss_Contra = self.criterion_Contra(features, labels)

        # KL divergence
        self.loss_KL = -torch.nn.functional.kl_div(self.f_Rss.softmax(dim=-1).log(), self.f_Rrs.softmax(dim=-1))

        # Consistency loss
        self.loss_Consis_s = self.criterion_MSE(self.pred_Rs+self.pred_Bs, self.Os)
        self.loss_Consis_r = self.criterion_MSE(self.pred_Rr+self.pred_Br, self.Or)
        self.loss_Consis_fBr = self.criterion_MSE(self.recon_f_Br, self.f_Br)
        self.loss_Consis_fBs = self.criterion_MSE(self.recon_f_Bs, self.f_Bs)
        # self.loss_Consis = self.loss_Consis_s + self.loss_Consis_r
        self.loss_Consis = self.loss_Consis_s + self.loss_Consis_r + self.loss_Consis_fBr + self.loss_Consis_fBs

        # Total loss
        loss_G = lambda_MSE*self.loss_MSE + self.loss_GAN + lambda_Contra * self.loss_Contra \
                 + lambda_Consis * self.loss_Consis + lambda_KL * self.loss_KL
        loss_G.backward()


    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # # forward
        self.forward()      # compute fake images and reconstruction images.

        # G
        self.set_requires_grad([self.netD_B], False)  # Ds require no gradients when optimizing Gs
        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero
        self.backward_G()             # calculate gradients for G_A and G_B
        self.optimizer_G.step()       # update G_A and G_B's weights

        # D_B
        self.set_requires_grad([self.netD_B], True)
        self.optimizer_D_B.zero_grad()
        self.backward_D_B()
        self.optimizer_D_B.step()
